
const WebSocket = require('ws');
const https = require('https');
// 创建一个 WebSocket 服务器，监听在指定端口
const wss = new WebSocket.Server({ port: 39001 });

const apiKey = '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';

// 当有客户端连接时执行
wss.on('connection', (ws) => {
    console.log('Client connected');

    let stopFlag = false

    function sendRequest(requestOptions, requestBody, retryCount = 3) {
        return new Promise((resolve, reject) => {
            const request = https.request(requestOptions, (response) => {

                response.setEncoding('utf8');
                ws.send(`openAIWSSStart`);

                response.on('data', (chunk) => {
                    if (stopFlag || !response.socket || !response.socket.readable) {
                        request.abort();
                        return;
                    }



                    let recMsg= JSON.parse(chunk.slice(6));

                    //console.log(recMsg.choices[0].delta.content);

                    //if(recMsg.choices[0].delta.finish_reason=='stop'){
                        //ws.send(`openAIWSSEnd`);
                        //request.abort();
                        //return
                    //}
        
                    // 发送消息给客户端
                    //ws.send(recMsg.choices[0].delta.content);
                });
    
                response.on('end', () => {
                    ws.send(`openAIWSSEnd`);
                    resolve({
                        code: 0,
                        data: 'openai end',
                    });
                });
            });
    
            request.on('abort', (res) => {
                console.log('abort request at ' + new Date()+`:${res}`);
                resolve({
                    code: 500,
                    data: '',
                });
            });
    
            request.on('error', (error) => {
                console.log('request error at ' + new Date()+`:${error}`);
                if (retryCount === 0) {
                    reject(error);
                } else {
                    retryCount--;
                    sendRequest(requestOptions, requestBody, retryCount).then(resolve).catch(reject);
                }
            });
    
            //发送消息
            request.write(requestBody);
            request.end();
        });
    }

    // 当收到消息时执行
    ws.on('message', (message) => {
        console.log(`Received: ${message}`);
        if(message=='connectOpenAIWSS'){
            ws.send(`connectOpenAIWSSBack`);
            return
        }

        let requestBody= message;
        let headers = {
            'Content-Type': 'application/json',
            'Authorization': 'Bearer ' + apiKey
        }

        let requestOptions = {
            'host': 'api.openai.com',
            'port': 443,
            'path': '/v1/chat/completions',
            'method': 'POST',
            'timeout': 1000000, //超时设置
            'headers': headers
        }

        ws.send(`openAIWSSBack`);

        let startTime = new Date().getTime();
        sendRequest(requestOptions, requestBody, 3)
            .then((res) => {
                let endTime = new Date().getTime();
                console.log('\r\n请求完成，耗时：' + (endTime - startTime) + '毫秒,' + `${requestBody}`);
            })
            .catch((error) => {
                let endTime = new Date().getTime();
                console.log('请求失败，耗时：' + (endTime - startTime) + '毫秒,' + `${requestBody}`);
                console.error(error);
            });

    });

    // 当客户端断开连接时执行
    ws.on('close', (res) => {
        console.log('Client disconnected:'+JSON.stringify(res));
        if(res===3333){
            stopFlag = true
        }
    });

    // 当客户端断开连接时执行
    ws.on('error', (res) => {
        console.log('wss error:'+JSON.stringify(res));
        stopFlag = true
    });
});